//////////////////////////////////////////////////////////////////
// (c) Copyright 2010-  by Jeongnim Kim
//////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////
//   National Center for Supercomputing Applications &
//   Materials Computation Center
//   University of Illinois, Urbana-Champaign
//   Urbana, IL 61801
//   e-mail: jnkim@ncsa.uiuc.edu
//
// Supported by 
//   National Center for Supercomputing Applications, UIUC
//   Materials Computation Center, UIUC
//////////////////////////////////////////////////////////////////
// -*- C++ -*-
/**@file pencildecomp.hpp
 */
#ifndef QMCPLUSPLUS_PENCILDECOMP_HPP
#define QMCPLUSPLUS_PENCILDECOMP_HPP
#include <config.h>
#include <string>
#include <iostream>
#include <iomanip>
#include <cstdlib>
#if (__GNUC__ == 2)
#include <strstream.h>
#else
#include <sstream>
#endif
#include <strstream>
#include <fstream>
#include <mpi/communicator.h>
#include <mpi/collectives.h>
#include <Utilities/Timer.h>
using namespace std;

namespace mpi
{

  template<typename T>
  struct PencilDecomp
  {
    /** enumeration for the data ordering */
    enum {PXYZ=0, PXZY, PYZX};
    /** enumerate row and column */
    enum {IROW=0, ICOL=1};
    /** enumerate the timer */
    enum {ROWTRANS_T=0 ,COLTRANS_T,NTIMERS=4};
    //const int Index[3][3];
    //enum {I_XYZ_X=0, I_XYZ_Y=1, I_XYZ_Z=2};
    //enum {I_XZY_X=0, I_XZY_Y=2, I_XZY_Z=1};
    //enum {I_YZX_X=2, I_YZX_Y=0, I_YZX_Z=1};
    /** size of the problem Length^3 */
    int Length;
    /** save the rank */
    int Rank;
    /** processor grid */
    int pGrid[2];
    /** grid coordination */
    int pCoord[2];
    /** lower bound of the block owned */
    int blockMin[2];
    /** upper bound of the block owned */
    int blockMax[2];
    /** size of the blocks */
    int blockSize[2];
    /** size of the messages per node */
    size_t msgSize[2];
    /** timers */
    double wtime[NTIMERS];
    /** comm for row */
    communicator *rowComm;
    /** comm for column */
    communicator *colComm;
    /** buffer to perform rowComm->alltoall */
    vector<T> rowBuffer;
    /** buffer to perform colComm->alltoall */
    vector<T> colBuffer;
    vector<int> rowCounts, rowOffsets;
    vector<int> colCounts, colOffsets;

    PencilDecomp(int l, int nx, int ny, int ipx, int ipy)
      :Length(l),Rank(ipy+ipx*ny),rowComm(0),colComm(0)
      {
        pGrid[0]=nx;pGrid[1]=ny;
        pCoord[0]=ipx;pCoord[1]=ipy;
        partition();
      }

    /** constructor
     * @param l Length
     * @param ny number of partitions in the column
     * @param c 
     */
    PencilDecomp(int l, int ny, communicator& c)
      :Length(l),rowComm(0),colComm(0)
      {
        Rank=c.rank();
        pGrid[0]=c.size()/ny;pGrid[1]=ny;
        pCoord[0]=Rank/ny;
        pCoord[1]=Rank%ny;
        setCommunicator(c);
        partition();
      }

    virtual ~PencilDecomp()
    {
      if(rowComm) delete rowComm;
      if(colComm) delete colComm;
    }

    inline void reset()
    {
      for(int i=0; i<NTIMERS; ++i) wtime[i]=0.0;
    }

    inline void partition()
    {
      partition(Length,pGrid[0],pCoord[0],blockSize,blockMin,blockMax,IROW);
      partition(Length,pGrid[1],pCoord[1],blockSize,blockMin,blockMax,ICOL);

      msgSize[IROW]=blockSize[IROW]*blockSize[ICOL]*Length/rowComm->size();
      msgSize[ICOL]=blockSize[IROW]*blockSize[ICOL]*Length/colComm->size();

      rowBuffer.resize(msgSize[IROW]*rowComm->size(),Rank);
      colBuffer.resize(msgSize[ICOL]*colComm->size(),Rank);

      rowCounts.resize(rowComm->size(),msgSize[IROW]);
      rowOffsets.resize(rowComm->size());
      for(int i=0; i<rowComm->size(); ++i) rowOffsets[i]=i*msgSize[IROW];

      colCounts.resize(colComm->size(),msgSize[ICOL]);
      colOffsets.resize(colComm->size());
      for(int i=0; i<colComm->size(); ++i) colOffsets[i]=i*msgSize[ICOL];
    }

    void setCommunicator(communicator& c)
    {
      if(rowComm) delete rowComm;
      if(colComm) delete colComm;
      rowComm= new communicator(c.split(pCoord[0],pCoord[1]));
      colComm= new communicator(c.split(pCoord[1],pCoord[0]));
    }

    void print(int plan, std::ostream& os)
    {
      char f[512];
      sprintf(f,"%6d %3d %3d %3d %3d %6d %6d %6d %6d %6d(%d) %6d(%d) "
          ,Rank, pCoord[0], pCoord[1], rowComm->rank(), colComm->rank()
          ,blockMin[IROW],blockMax[IROW]
          ,blockMin[ICOL],blockMax[ICOL]
          ,blockMax[IROW]-blockMin[IROW] ,blockSize[IROW]
          ,blockMax[ICOL]-blockMin[ICOL] ,blockSize[ICOL]
          );
      os << f;
    }

    inline void fft1d(int dir)
    {
#if defined(DOSTUFF)
      if(dir)
        spectrum();
      else
        rk4();
#endif
    }

    inline void spectrum()
    {
#if defined(DOSTUFF)
      std::copy(rowBuffer.begin(),rowBuffer.end(),colBuffer.begin());
#endif
    }

    inline void rk4()
    {
#if defined(DOSTUFF)
      std::copy(colBuffer.begin(),colBuffer.end(),rowBuffer.begin());
#endif
    }


    void print(std::ostream& os)
    {
      os << "XYZ "; print(PXYZ,os); 
      os << " size= " << rowBuffer.size() 
        << " msgSize " << msgSize[IROW]*sizeof(T) << " " << msgSize[ICOL]*sizeof(T) << endl;
      //os << "XZY "; print(P_XZY,os);
      //os << "YZX "; print(P_YZX,os);
      /*
      const int nb=2;
      {
        std::vector<double> v(rowComm->size()*nb);
        std::vector<double> vout(rowComm->size()*nb,-1);
        for(int i=0,ii=0; i< rowComm->size(); ++i)
          for(int k=0;k<nb;++k, ++ii) v[ii]=Rank*nb+k;
        alltoall(*rowComm,&v[0],&vout[0],nb);

        os << "Row-In ";
        for(int i=0; i<v.size(); ++i) os << setw(4) << v[i];
        os << endl;

        os << "Row-Out ";
        for(int i=0; i<vout.size(); ++i) os << setw(4) << vout[i];
        os << endl;
      }

      {
        std::vector<double> v(colComm->size()*nb);
        std::vector<double> vout(colComm->size()*nb,-1);
        for(int i=0,ii=0; i< colComm->size(); ++i)
          for(int k=0;k<nb;++k, ++ii) v[ii]=Rank*nb+k;
        alltoall(*colComm,&v[0],&vout[0],nb);

        os << "Col-In ";
        for(int i=0; i<v.size(); ++i) os << setw(4) << v[i];
        os << endl;

        os << "Col-Out ";
        for(int i=0; i<vout.size(); ++i) os << setw(4) << vout[i];
        os << endl;
      }
      */
      //os << "Reduce " << v << " " << vr << " " << vc << std::endl;
    }

    /** partition a vector n by p
     * @param n dimension in 1D
     * @param p number of partitions
     * @param ip rank in p
     * @param bsize size of the block (n%p)? n/p+1 :n/p;
     * @param bmin lower bound
     * @param bmax upper bound
     * @param index for the plan
     */
    inline void partition(int n, int p, int ip, int *bsize, int *bmin, int *bmax, int index)
    {
      int np=n/p;//n per partition
      int p_more=n-p*np;//partitions n/p
      int p_less=p-p_more;//partitions with n/p+1
      bsize[index]=(p_more)?np+1:np;//size
      if(ip<p_less)
      {
        bmin[index]=np*ip;
        bmax[index]=bmin[index]+np;
      }
      else
      {
        bmin[index]=(np+1)*(ip-p_less)+np*p_less;
        bmax[index]=bmin[index]+np+1;
      }
    }
  };
}
#endif


/***************************************************************************
 * $RCSfile$   $Author: jnkim $
 * $Revision: 1770 $   $Date: 2007-02-17 17:45:38 -0600 (Sat, 17 Feb 2007) $
 * $Id: fft2d.cpp 1770 2007-02-17 23:45:38Z jnkim $ 
 ***************************************************************************/
